#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2020/1/17 下午3:24
# @Author  : fugang_le
# @Software: PyCharm

import numpy as np
import jieba
from gensim.models import Word2Vec

MAX_WORDS_IN_BATCH = 10000
stopwords = []


# 对每个句子的所有词向量取均值
def buildWordVector(text, imdb_w2v):
    words = cw(text)
    if not words:
        return
    size = imdb_w2v.size
    vec = np.zeros(size).reshape((1, size))
    count = 0.
    for word in text:
        try:
            vec += imdb_w2v[word].reshape((1, size))
            count += 1.
        except KeyError:
            continue
    if count != 0:
        vec /= count
    return vec



class W2V():

    def __init__(self,sentences=None, corpus_file=None, size=100, alpha=0.025, window=5, min_count=5,
                 max_vocab_size=None, sample=1e-3, seed=1, workers=3, min_alpha=0.0001,
                 sg=0, hs=0, negative=5, ns_exponent=0.75, cbow_mean=1, hashfxn=hash, iter=5, null_word=0,
                 trim_rule=None, sorted_vocab=1, batch_words=MAX_WORDS_IN_BATCH, compute_loss=False, callbacks=(),
                 max_final_vocab=None, **kwargs):
        self.model = None
        self.parameter = {}
        self.parameter['sentences'] = sentences
        self.parameter['corpus_file'] = corpus_file
        self.parameter['size'] = size
        self.parameter['alpha'] = alpha
        self.parameter['window'] = window
        self.parameter['min_count'] = min_count
        self.parameter['max_vocab_size'] = max_vocab_size
        self.parameter['sample'] = sample
        self.parameter['seed'] = seed
        self.parameter['workers'] = workers
        self.parameter['min_alpha'] = min_alpha
        self.parameter.update(kwargs)

    def train(self,):
        self.model = Word2Vec(**self.parameter)

    def save(self, model_path):
        self.model.save(model_path)


if __name__ == '__main__':
    model_path = 'w2v_model.pkl'
    sentences = [['first', 'sentence'], ['second', 'sentence']]
    model = W2V(sentences=sentences, min_count=1, )
    model.train()
    print(model.model['first'])
    model.save(model_path)